import numpy as np
import plotly.graph_objs as go

N = 1 + 100

# Define the function f(t) for a given n
def f(t, n):    return 4/np.pi*sum( np.sin((2*i+1)*t) /(2*i+1) for i in range(0, n+1)) ### Square Wave
def f(t, n):    return 1/2 - 1/np.pi*sum( np.sin((i)*t) /(i) for i in range(1, n+1)) ### Sawtooth v1
def f(t, n):    return 2/np.pi*sum( np.sin((i)*t) /(i) for i in range(1, n+1)) ### Sawtooth v2
def f(t, n):    return 8*(1/np.pi)**2*sum( np.sin((2*i+1)*t)*(-1)**i/(2*i+1)**2 for i in range(0, n+1)) ### Triangle
def f(t, n):    return sum(np.sin(i**2 * t) / i**1 for i in range(1, n+1)) # Riemann pathological
def f(t, n):    return sum(np.sin(i**2 * t) / i**2 for i in range(1, n+1)) # Riemann pathological

# Generate t values
t_values = np.linspace(0, 2*np.pi, 1280)

# Create frames for the animation
frames = []
for n in range(1, N,10):  # Start from n=2 since n=1 is static
    f_values = [f(t, n) for t in t_values]
    frames.append(go.Frame(
        data=[
            go.Scatter(x=t_values, y=[f(t, 1) for t in t_values], mode='lines', name='n=1 (dashed)', line=dict(dash='dash')),
            go.Scatter(x=t_values, y=[f(t, 3) for t in t_values], mode='lines', name='n=3 (dashed)', line=dict(dash='dash')),
            go.Scatter(x=t_values, y=f_values, mode='lines', name=f'n={n}'),
        ],
        name=str(n),
        layout=go.Layout(title=f'Evolution of f(t) as n = {n}')
    ))

# Create the initial plot
fig = go.Figure(
    data=[
        go.Scatter(x=t_values, y=[f(t, 1) for t in t_values], mode='lines', name='n=1', line=dict(dash='dash')),
        go.Scatter(x=t_values, y=[f(t, 3) for t in t_values], mode='lines', name='n=3', line=dict(dash='dash')),
        go.Scatter(x=t_values, y=f_values, mode='lines', name=f'n={n}'),
    ],
    layout=go.Layout(
        title=f'Evolution of f(t) = sum(sin(n^2 * t) / n^2) as n increases from 1 to {N-1}',
        xaxis_title='t',
        yaxis_title='f(t)',
        plot_bgcolor='white',
        updatemenus=[{
            'buttons': [
                {
                    'args': [None, {'frame': {'duration': 500, 'redraw': True}, 'fromcurrent': True}],
                    'label': 'Play',
                    'method': 'animate'
                },
                {
                    'args': [[None], {'frame': {'duration': 0, 'redraw': True}, 'mode': 'immediate', 'transition': {'duration': 0}}],
                    'label': 'Pause',
                    'method': 'animate'
                }
            ],
            'type': 'buttons'
        }]
    ),
    frames=frames
)
fig.update_layout(width=1000, height=600)
fig.show()